from AStock.ASSetting import settings
from AStock.ASUtil import date_to_timestamp
from AStock.ASQuery import ASQuery_nearest_trade_date


class ASIndexChangeRatio(object):
    def __init__(self, begin_date, end_date, index=None):
        """
        # 默认获取的指数：
        # 上证指数：000001
        # 创业板指：399006
        # 沪深300：000300
        # 上证50：000016
        # 中证500：000905
        # 创业板综：399102
        # 中证1000：000852
        :param index: str or list of index code
        """
        self._begin_date = begin_date
        self._end_date = end_date

        self._index_dict = {
            '000001': '上证指数',
            '399006': '创业板指',
            '000300': '沪深300',
            '000016': '上证50',
            '000905': '中证500',
            '399102': '创业板综',
            '000852': '中证1000'
        }
        if index is None:
            self._index = [
                '000001',
                '399006',
                '000300',
                '000016',
                '000905',
                '399102',
                '000852'
            ]
        elif isinstance(index, str):
            self._index = [index]
        elif isinstance(index, list):
            self._index = index
        else:
            raise TypeError('index must be str/list/None')

    def run(self):
        begin_date = ASQuery_nearest_trade_date(self._begin_date)
        end_date = ASQuery_nearest_trade_date(self._end_date)
        if begin_date is None:
            raise ValueError('cannot find nearest trade date of begin date')
        if end_date is None:
            raise ValueError('cannot find nearest trade date of end date')

        begin_date_stamp = date_to_timestamp(begin_date)
        end_date_stamp = date_to_timestamp(end_date)
        begin_data = self._query_close_points(begin_date_stamp)
        end_data = self._query_close_points(end_date_stamp)

        for code_item, begin_data_item in begin_data.items():
            if code_item not in end_data:
                print('{}\t{}\t{}'.format(code_item, self._index_dict.get(code_item), '--'))
            else:
                end_data_item = end_data[code_item]
                # print(begin_data_item)
                # print(end_data_item)
                ratio = (end_data_item['close'] - begin_data_item['close']) / begin_data_item['close']
                ratio = round(ratio * 100, 2)
                print('{}\t{}\t{}%'.format(code_item, self._index_dict.get(code_item), ratio))

    def _query_close_points(self, date_stamp):
        coll = settings.database.index_day
        filter_ = {
            'code': {'$in': self._index},
            'date_stamp': date_stamp
        }
        projection = {'_id': 0, 'code': 1, 'date': 1, 'close': 1}
        try:
            cursor = coll.find(filter_, projection)
        except:
            raise

        data = {}
        for row in cursor:
            data.__setitem__(row['code'], row)

        return data


